package tcp_proxy_middleware

import (
	"fmt"
	"strings"

	"gin-gateway/core/utils"
	"gin-gateway/reverse_proxy"
	"log"
)

// TCPWhiteAndBlackListMiddleware 黑/白名单
func TCPWhiteAndBlackListMiddleware() func(c *reverse_proxy.TcpSliceRouterContext) {
	return func(c *reverse_proxy.TcpSliceRouterContext) {
		serviceDetail := GetServiceDetail(c)

		clientIP := strings.Split(c.Conn.RemoteAddr().String(), ":")[0]

		var ipBlackList []string
		if serviceDetail.AccessControl.BlackList != "" {
			ipBlackList = strings.Split(serviceDetail.AccessControl.BlackList, ",")
		}
		if serviceDetail.AccessControl.OpenAuth == 1 && len(ipBlackList) > 0 && utils.StrInSlice(ipBlackList, clientIP) {
			_, err := c.Conn.Write([]byte(fmt.Sprintf("%s in ip black list", clientIP)))
			if err != nil {
				log.Println("TCPBlackListMiddleware Write error:", err)
			}
			c.Abort()
			return
		}

		var ipWhiteList []string
		if serviceDetail.AccessControl.WhiteList != "" {
			ipWhiteList = strings.Split(serviceDetail.AccessControl.WhiteList, ",")
		}
		if serviceDetail.AccessControl.OpenAuth == 1 && len(ipWhiteList) > 0 && !utils.StrInWhileList(ipWhiteList, clientIP) {
			_, err := c.Conn.Write([]byte(fmt.Sprintf("%s not in ip white list", clientIP)))
			if err != nil {
				log.Println("TCPBlackListMiddleware Write error:", err)
			}
			c.Abort()
			return
		}
		c.Next()
	}
}
